package file

import (
	"fmt"
	"os"
	"sync"

	"github.com/boltdb/bolt"
)

// BoltDB bolt db
type BoltDB struct {
	File   string
	Bucket string
	db     *bolt.DB
	sync.Mutex
}

// NewBoltDB 新建bolt db
func NewBoltDB(file string) *BoltDB {
	return &BoltDB{
		File:   file,
		Bucket: "default",
		db:     nil,
	}
}

// open 打开文件
func (r *BoltDB) open() error {
	r.Lock()
	defer r.Unlock()
	if r.db == nil {
		db, err := bolt.Open(r.File, os.ModePerm, nil)
		if err != nil {
			return err
		}
		err = db.Update(func(tx *bolt.Tx) error {
			_, err := tx.CreateBucketIfNotExists([]byte(r.Bucket))
			return err
		})
		if err != nil {
			return err
		}
		r.db = db
	}
	return nil
}

// Write 写入数据
func (r *BoltDB) Write(data chan map[string]string) error {
	if err := r.open(); err != nil {
		return err
	}
	return r.db.Update(func(tx *bolt.Tx) error {
		b, err := tx.CreateBucketIfNotExists([]byte(r.Bucket))
		if err != nil {
			return err
		}
		for value := range data {
			for k, v := range value {
				if err := b.Put([]byte(k), []byte(v)); err != nil {
					return err
				}
			}
		}
		return nil
	})
}

// Range 遍历数据
func (r *BoltDB) Range(handle func(key, value string) error) error {
	if err := r.open(); err != nil {
		return err
	}
	return r.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(r.Bucket))
		if b == nil {
			return fmt.Errorf("bucket '%s' is not exist", r.Bucket)
		}
		c := b.Cursor()
		for k, v := c.First(); k != nil; k, v = c.Next() {
			if err := handle(string(k), string(v)); err != nil {
				return err
			}
		}
		return nil
	})
}

// Find 查找数据
func (r *BoltDB) Find(input, output chan string) error {
	if err := r.open(); err != nil {
		return err
	}
	return r.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(r.Bucket))
		if b == nil {
			return fmt.Errorf("bucket '%s' is not exist", r.Bucket)
		}
		for key := range input {
			v := b.Get([]byte(key))
			output <- string(v)
		}
		return nil
	})
}

// Remove 移除KEY
func (r *BoltDB) Remove(keys ...string) error {
	if err := r.open(); err != nil {
		return err
	}
	return r.db.Update(func(tx *bolt.Tx) error {
		if b := tx.Bucket([]byte(r.Bucket)); b != nil {
			for _, key := range keys {
				if err := b.Delete([]byte(key)); err != nil {
					return err
				}
			}
		}
		return nil
	})
}

// Close 关闭文件
func (r *BoltDB) Close() error {
	r.Lock()
	defer r.Unlock()
	if r.db != nil {
		err := r.db.Close()
		r.db = nil
		return err
	}
	return nil
}
